import os  
import torch
import torchvision
from utils.utils import get_accurary
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader 
import numpy as np
from torchvision import transforms
from PIL import Image  
from dataset.MultiLabelImageDatasetTest import MultiLabelImageDatasetTest
from datetime import datetime, timedelta
from utils.utils import set_seed
from config import opt
from scipy import spatial

set_seed(opt.seed)



net = torch.load(opt.model_path)
net = net.to(opt.device)
net.eval()





transformer = transforms.Compose([
    
    transforms.Resize(256),
    transforms.CenterCrop(224),  # 通常在调整大小后进行中心裁剪

    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

#获取训练集已知信号异常分数
train_dataset = MultiLabelImageDatasetTest(opt.path + 'train/', transform=transformer)
train_loader = DataLoader(train_dataset, batch_size = opt.test_batch_size, shuffle=False)
train_dis_list = [] 
with torch.no_grad():
    for i, data_info in enumerate(train_loader):
        images, labels = data_info['image_tensors'].to(opt.device), data_info['labels']
        outputs,trans_feature,features = net(images)
        features = features.detach().cpu().numpy()

        ave_trans_feature = np.zeros_like(features)
        for i in trans_feature.keys():
            ave_trans_feature  += trans_feature[i][0].detach().cpu().numpy()
        ave_trans_feature = ave_trans_feature/3
        dis =  1 - spatial.distance.cosine(ave_trans_feature[0] , features[0])
        train_dis_list.append(dis)






# 获取test异常分数
test_dataset = MultiLabelImageDatasetTest(opt.path + 'test/', transform=transformer)
test_loader = DataLoader(test_dataset, batch_size = opt.test_batch_size, shuffle=False)
test_dis_list = [] 
detect_res = []
ground_truth = []
with torch.no_grad():
    for i, data_info in enumerate(test_loader):
        images, labels = data_info['image_tensors'].to(opt.device), data_info['labels']
        if(labels['T000'].item() == 1):
            ground_truth.append(1)
        else:
            ground_truth.append(0)

        outputs,trans_feature,features = net(images)
        features = features.detach().cpu().numpy()

        ave_trans_feature = np.zeros_like(features)
        for i in trans_feature.keys():
            ave_trans_feature  += trans_feature[i][0].detach().cpu().numpy()
        ave_trans_feature = ave_trans_feature/3
        dis =  1 - spatial.distance.cosine(ave_trans_feature[0] , features[0])
        test_dis_list.append(dis)

plt.hist(test_dis_list, bins=50, color="red", alpha=.9)
plt.hist(train_dis_list, bins=50, color="blue", alpha=.9)
plt.savefig("./picture/dis" + opt.log_time +"_pred.png") # 你给图片起的名字


